from enum import Enum
from typing import Optional, List, Dict, Any
from pydantic import BaseModel
from app.schemas.facenet_base_schema import TaskStatus
# 1. 任务状态
class TaskStatus(str, Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"
# 演化预测算法的输入参数
class EvolutionPredictInputParams(BaseModel):
    historical_data: List[List[List[int]]]  # 历史时间步的数据，格式为[[[node1,node2,weight], ...], [...]]
    prediction_steps: int = 3  # 预测未来的时间步数
    window_size: int = 5  # 用于预测的历史窗口大小
    node_count: int = 128  # 网络中的节点数量
    prediction_method: str = "linear"  # 预测方法，可以是 "linear", "exp", "lstm" 等

# 每个时间步的预测结果
class TimePredictionResult(BaseModel):
    time_step: int  # 时间步
    predicted_edges: List[List[int]]  # 预测的边列表 [[node1,node2,weight], ...]
    confidence_score: float  # 预测的置信度
    prediction_metrics: Dict[str, float]  # 预测的评估指标

# 演化预测算法的输出参数
class EvolutionPredictOutputParams(BaseModel):
    predictions: List[TimePredictionResult]  # 每个预测时间步的结果
    algorithm: str = "EvolutionPredict"  # 算法名称
    parameters: Dict[str, Any]  # 算法参数
    overall_accuracy: float  # 整体预测准确率

# 在现有的InputParams和OutputParams中添加新的字段
class InputParams(BaseModel):
    evolution_predict_params: Optional[EvolutionPredictInputParams] = None
    # ... 其他现有参数 ...

class OutputParams(BaseModel):
    evolution_predict_results: Optional[EvolutionPredictOutputParams] = None
    # ... 其他现有参数 ... 


#5. 算法请求
class AlgorithmRequest(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    input_params: InputParams

# 6. 算法中间响应
class AlgorithmMiddleResponse(BaseModel):
    task_id: str
    task_callback_url: str
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None

    input_params: InputParams
    error_message: Optional[str] = None


# 7. 算法最终响应
class AlgorithmResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None

    error_message: Optional[str] = None

    output_params: OutputParams

